import screeninfo
import win32api
import win32con
import win32gui
from collections.abc import Callable
import numpy as np
import os
import cv2


DEFAULT_WORK_DIR = os.path.dirname(__file__)


class Screen(screeninfo.Monitor):
    def __init__(self, monitor: screeninfo.Monitor, alias: str, wallpaper: np.ndarray = None, contain_work_screen: bool = True) -> None:
        super(Screen, self).__init__(**monitor.__dict__)
        self.alias: str = alias
        self.contain_work_screen = contain_work_screen
        self.__get_work_screen()
        self.__calculate_geometry()
        # 几何尺寸
        self.work: Screen | None
        self.work_shift_x: int
        self.work_shift_y: int
        self.start_center: tuple[float, float]
        self.center: tuple[float, float]
        self.end_center: tuple[float, float]
        # 壁纸图像
        self.wallpaper: np.ndarray | None = wallpaper

    def setWallpaper(self, image: np.ndarray) -> None:
        self.wallpaper = image

    def __repr__(self) -> str:
        string = "Screen(\n"
        for key, value in self.__dict__.items():
            if value is not None:
                value_str = repr(value) if type(value) not in [
                    np.ndarray] else str(type(value))
                string += f"\t{key}={value_str}\n"
        string += ")\n"
        return string

    def __get_work_screen(self) -> None:
        if self.contain_work_screen == True:
            monitor_list = win32api.EnumDisplayMonitors()
            monitor_list = map(
                lambda m: win32api.GetMonitorInfo(m[0]), monitor_list)
            monitor_list = list(
                filter(lambda m: m['Device'] == self.name, monitor_list))
            assert len(
                monitor_list) == 1, f"matched {len(monitor_list)} screen(s) named '{self.name}'"
            monitor = monitor_list[0]
            work = screeninfo.Monitor(
                x=monitor['Work'][0],
                y=monitor['Work'][1],
                width=monitor['Work'][2] - monitor['Work'][0],
                height=monitor['Work'][3] - monitor['Work'][1],
                width_mm=None,
                height_mm=None,
                name=self.name+':work',
                is_primary=None
            )
            self.work = Screen(work, self.alias+':work',
                               contain_work_screen=False)

    def __calculate_geometry(self) -> None:
        if self.contain_work_screen == True:
            self.work_shift_x = self.work.x - self.x
            self.work_shift_y = self.work.y - self.y
            if self.work.width > self.work.height:
                self.start_center = (
                    (self.work_shift_x + self.work.height/2) / self.width,
                    (self.work_shift_y + self.work.height/2) / self.height,
                )
                self.end_center = (
                    (self.work_shift_x + self.work.width -
                     self.work.height/2) / self.width,
                    (self.work_shift_y + self.work.height/2) / self.height,
                )
            else:
                self.start_center = (
                    (self.work_shift_x + self.work.width/2) / self.width,
                    (self.work_shift_y + self.work.width/2) / self.height,
                )
                self.end_center = (
                    (self.work_shift_x + self.work.width/2) / self.width,
                    (self.work_shift_y + self.work.height -
                     self.work.width/2) / self.height,
                )
            self.center = (
                (self.work_shift_x + self.work.width/2) / self.width,
                (self.work_shift_y + self.work.height/2) / self.height,
            )


class DisplayManager():

    def __init__(self, work_dir: str = DEFAULT_WORK_DIR) -> None:
        self.work_dir: str = work_dir
        self.wp_save_path: str = os.path.join(self.work_dir, 'wallpaper.jpg')
        monitor_list = screeninfo.get_monitors()
        screen_list = []
        for i, monitor in enumerate(monitor_list):
            screen_list.append(Screen(monitor, f'Screen-{i+1}'))
        self.screen_list = screen_list
        self.__merge_screens()
        self.virtual_screen: Screen
        self.default_wallpaper_painter: Callable[[Screen], np.ndarray]

    def full_update(self):
        self.__create_default_wallpaper_for_empty_screens()
        self.__merge_all_wallpapers()
        self.__set_wallpaper()

    def update_screen(self, screen: str) -> None:
        if type(screen) == str:
            screen = self[screen]

    def __getitem__(self, __name: str) -> Screen | list[Screen] | None:
        res = list(filter(lambda d: d.alias == __name, self.screen_list))
        assert len(res) <= 1, f"found more than one screen named {__name}"
        if len(res) == 1:
            return res[0]
        else:
            return None

    def __iter__(self):
        for screen in self.screen_list:
            screen: Screen
            yield screen

    def __contains__(self, alias) -> bool:
        return self[alias] != None

    def __repr__(self) -> str:
        return repr(self.screen_list)

    def __set_wallpaper(self) -> None:
        cv2.imwrite(self.wp_save_path, self.virtual_screen.wallpaper, [
                    int(cv2.IMWRITE_JPEG_QUALITY), 100])
        key = win32api.RegOpenKeyEx(win32con.HKEY_CURRENT_USER,
                                    "Control Panel\\Desktop", 0, win32con.KEY_SET_VALUE)
        win32api.RegSetValueEx(key, "WallpaperStyle", 0,
                               win32con.REG_SZ, "10")  # 设置壁纸的布置方式
        win32api.RegSetValueEx(key, "TileWallpaper", 0, win32con.REG_SZ, "1")
        win32gui.SystemParametersInfo(
            win32con.SPI_SETDESKWALLPAPER, self.wp_save_path, win32con.SPIF_SENDWININICHANGE)

    def __create_default_wallpaper_for_empty_screens(self) -> None:
        for screen in self:
            if screen.wallpaper is None:
                screen.wallpaper = self.default_wallpaper_painter(screen)

    def __merge_screens(self) -> None:
        left = min(map(lambda d: d.x, self))
        right = max(map(lambda d: d.x + d.width, self))
        top = min(map(lambda d: d.y, self))
        bottom = max(map(lambda d: d.y + d.height, self))
        width = right - left
        height = bottom - top
        self.virtual_screen = Screen(
            monitor=screeninfo.Monitor(left, top, width, height),
            alias='main',
            contain_work_screen=False
        )

    def __merge_all_wallpapers(self) -> None:
        if self.virtual_screen.wallpaper is None:
            self.virtual_screen.wallpaper = np.empty(
                (self.virtual_screen.height, self.virtual_screen.width, 3), dtype=np.uint8)
        for screen in self:
            if screen.wallpaper is not None:
                left = screen.x - self.virtual_screen.x
                right = left + screen.width
                top = screen.y - self.virtual_screen.y
                bottom = top + screen.height
                self.virtual_screen.wallpaper[top:bottom,
                                              left:right] = screen.wallpaper
